import cv2
import numpy as np
from fastapi import FastAPI, File, UploadFile, Form, HTTPException
from fastapi.responses import JSONResponse
import uvicorn

app = FastAPI()

# 加载相机标定参数（需要根据实际路径修改）
try:
    with np.load('camera_params.npz') as X:
        mtx, dist = [X[i] for i in ('mtx', 'dist')]
except Exception as e:
    raise RuntimeError(f"无法加载相机标定参数: {str(e)}")

# 定义ArUco字典和检测器
aruco_dict = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_1000)
detector = cv2.aruco.ArucoDetector(aruco_dict)


def calculate_corner_points(rvec, tvec, boxW, boxFH, boxBh, boxD):
    """
    计算料盒的八个角点在相机坐标系中的位置
    :param rvec: ArUco码的旋转向量
    :param tvec: ArUco码的平移向量
    :param boxW: 料盒宽度（米）
    :param boxFH: 料盒前高（米）
    :param boxBh: 料盒后高（米）
    :param boxD: 料盒深度（米）
    :return: 料盒角点坐标和ArUco坐标
    """
    try:
        # ArUco码中心到底边的距离（根据实际情况调整）
        aruco_offset_y = 0.058  # 5.8cm

        # 定义料盒角点相对坐标
        halfW = boxW / 2
        halfD = boxD / 2

        local_corners = np.array([
            # 底面
            [-halfW, -aruco_offset_y, 0],  # 左后下
            [halfW, -aruco_offset_y, 0],  # 右后下
            [halfW, boxFH - aruco_offset_y, 0],  # 右前下
            [-halfW, boxFH - aruco_offset_y, 0],  # 左前下
            # 顶面后部
            [-halfW, -aruco_offset_y, -boxD],  # 左后上
            [halfW, -aruco_offset_y, -boxD],  # 右后上
            # 顶面前部
            [halfW, boxBh - aruco_offset_y, -boxD],  # 右前上
            [-halfW, boxBh - aruco_offset_y, -boxD],  # 左前上
        ], dtype=np.float32)

        # 转换到相机坐标系
        rotation_matrix, _ = cv2.Rodrigues(rvec)
        box_corners = (rotation_matrix @ local_corners.T).T + tvec.reshape(1, 3)

        return {
            "aruco_position": tvec.tolist()[0],
            "box_corners": box_corners.tolist()
        }
    except Exception as e:
        raise RuntimeError(f"坐标计算失败: {str(e)}")


@app.post("/process")
async def process_box_image(
        aruco_id: int = Form(...),
        boxW: float = Form(...),
        boxFH: float = Form(...),
        boxBh: float = Form(...),
        boxD: float = Form(...),
        image: UploadFile = File(...)
):
    """
    处理料盒图像API接口
    参数：
    - aruco_id: 要处理的ArUco码ID
    - boxW: 料盒宽度（米）
    - boxFH: 料盒前高（米）
    - boxBh: 料盒后高（米）
    - boxD: 料盒深度（米）
    - image: 上传的图像文件
    """
    try:
        # 读取并解码图像
        img_data = await image.read()
        frame = cv2.imdecode(np.frombuffer(img_data, np.uint8), cv2.IMREAD_COLOR)

        # 检测ArUco码
        corners, ids, _ = detector.detectMarkers(frame)
        if ids is None:
            raise HTTPException(status_code=400, detail="未检测到ArUco码")

        # 查找指定ID的ArUco码
        target_idx = None
        for i, marker_id in enumerate(ids.flatten()):
            if marker_id == aruco_id:
                target_idx = i
                break

        if target_idx is None:
            raise HTTPException(status_code=404, detail=f"未找到ID为{aruco_id}的ArUco码")

        # 估计姿态
        rvecs, tvecs, _ = cv2.aruco.estimatePoseSingleMarkers(
            corners[target_idx],
            markerLength=0.043,  # ArUco码实际尺寸（单位：米）
            cameraMatrix=mtx,
            distCoeffs=dist
        )

        # 计算角点坐标
        result = calculate_corner_points(
            rvecs[0], tvecs[0],
            boxW=boxW,
            boxFH=boxFH,
            boxBh=boxBh,
            boxD=boxD
        )

        return JSONResponse(content={
            "status": "success",
            "aruco_id": aruco_id,
            "aruco_position": result["aruco_position"],
            "box_corners": result["box_corners"]
        })

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)